from sklearn.model_selection import KFold, train_test_split
import numpy as np
import pandas as pd
import math

class Data_Helper:

    def fixed_data(self, x_data, y_data, fixed_size=600):
        '''
        对于小于这个长度的数据进行padding，对于大于这个长度的数据进行截取
        :param x_data:  x数据
        :param y_data:  y数据
        :param fixed_size: 截取大小
        :return: 返回x和y的数据，并且为numpy类型
        '''
        for i in range(len(x_data)):
            t = x_data[i]
            if len(t) > fixed_size:
                x_data[i] = t[:fixed_size]
            while len(t) < fixed_size:
                t.append(-1)
        x_data = np.array(x_data)
        y_data = np.array(y_data)
        return x_data, y_data

    def KFold(self, x_data, y_data, K=10, shuffle=True):
        '''
        默认为10倍交叉验证，此函数为初始化
        :param x_data: x数据
        :param y_data: y数据
        :param K: 几倍交叉
        :return: 返回训练数据，测试数据
        '''
        res = []
        kf = KFold(n_splits=K, shuffle=shuffle)
        for train_index, test_index in kf.split(x_data):
            x_train = x_data[train_index]
            x_test = x_data[test_index]
            y_train = y_data[train_index]
            y_test = y_data[test_index]
            res.append([x_train, x_test, y_train, y_test])
        return res

    def train_test_split(self, x_data, y_data, test_size=0.3, shuffle=True):
        return train_test_split(x_data, y_data, test_size=test_size, shuffle=shuffle)

    def split_user_data(self, x_data, y_data, z_data):
        res = {}
        for i in range(len(z_data)):
            filename = z_data[i][z_data[i].rfind('/') + 1:]
            user = int(ord(filename.split('_')[1]) - ord('a')) + 1
            if user in res:
                res[user]['x_data'].append(x_data[i])
                res[user]['y_data'].append(y_data[i])
            else:
                res[user] = {
                    'x_data': [],
                    'y_data': []
                }
        return res
